#!/bin/bash
# 支持sda2、mmcblk1p2、nvme0n1p3格式
# 如果工具包没有安装，请先安装sudo apt update && sudo apt install cloud-guest-utils

# 检测root
if ! [ $(id -u) = 0 ]; then
  echo "Please run as root"
  exit 1
fi

# 检测工具包
if ! [ -x "$(command -v growpart)" ]; then
  echo "growpart not found, please install cloud-guest-utils"
  echo "sudo apt update && sudo apt install cloud-guest-utils"
  exit 1
fi

ROOT_DEV=$(findmnt -n -o SOURCE /)

# 处理分区命名差异
if [[ $ROOT_DEV =~ ^/dev/[a-z]+[0-9]+$ ]]; then
  # 处理 sda2, sdb2 格式
  DEVICE_NAME=$(echo $ROOT_DEV | sed 's/[0-9]*$//')
  PART_NUM=$(echo $ROOT_DEV | sed 's/.*[^0-9]//')
elif [[ $ROOT_DEV =~ ^/dev/[a-z]+[0-9]+p[0-9]+$ ]]; then
  # 处理 mmcblk1p2, nvme0n1p3 格式
  DEVICE_NAME=$(echo $ROOT_DEV | sed 's/p[0-9]*$//')
  PART_NUM=$(echo $ROOT_DEV | sed 's/.*p//')
else
  echo "Unsupported partition format: $ROOT_DEV" >&2
  exit 1
fi

echo "Detected root device: $ROOT_DEV"
echo "Using device: $DEVICE_NAME, partition: $PART_NUM"

if growpart $DEVICE_NAME $PART_NUM ; then
  case $(blkid -s TYPE -o value $ROOT_DEV) in
    ext4)  resize2fs $ROOT_DEV ;;
    xfs)   xfs_growfs / ;;
    btrfs) btrfs filesystem resize max / ;;
    *)     echo "Unsupported FS: $(blkid -s TYPE -o value $ROOT_DEV)" >&2; exit 1 ;;
  esac
fi